import { ref, computed, watch } from 'vue'
import type { RowData } from './types'

interface FlatRow {
  /** 原始行对象 */
  row: RowData
  /** 树层级深度（0 为根） */
  depth: number
}

export interface UseTreeRowsReturn {
  /** 扁平可渲染数组 */
  flatRows: ReturnType<typeof computed<FlatRow[]>>
  collapse: (id: string|number) => void
  expand: (id: string|number) => void
  toggle: (id: string|number) => void
  isCollapsed: (id: string|number) => boolean
}

/** DFS 生成扁平数组（尊重折叠状态） */
function flatten(
  nodes: RowData[],
  collapsed: Set<string|number>,
  depth = 0,
  acc: FlatRow[] = []
) {
  for (const n of nodes) {
    acc.push({ row: n, depth })
    if (n.children?.length && !collapsed.has(n.id)) {
      flatten(n.children, collapsed, depth + 1, acc)
    }
  }
  return acc
}

/**
 * 树形行折叠逻辑
 * @param treeData 响应式原始树
 * @param defaultExpandedKeys 初始展开 id
 */
export function useTreeRows(
  treeData: RowData[],
  defaultExpandedKeys: (string|number)[]
): UseTreeRowsReturn {
  /* 折叠集合：默认折叠 -> 所有有 children 的行；再减去默认展开 id */
  const collapsed = ref<Set<string|number>>(new Set())

  const init = (nodes: RowData[]) => {
    const walk = (arr: RowData[]) => {
      arr.forEach(n => {
        if (n.children?.length) collapsed.value.add(n.id)
        if (n.children?.length) walk(n.children)
      })
    }
    walk(nodes)
    defaultExpandedKeys.forEach(id => collapsed.value.delete(id))
  }
  init(treeData)

  /** 计算扁平数组 */
  const flatRows = computed<FlatRow[]>(() =>
    flatten(treeData, collapsed.value)
  )

  /* 数据源变化时同步折叠集合 */
  watch(
    () => JSON.stringify(treeData.map(n => n.id)),  // 粗略依赖
    () => {
      collapsed.value.clear()
      init(treeData)
    }
  )

  /* 操作方法 */
  const collapse = (id: string|number) => collapsed.value.add(id)
  const expand   = (id: string|number) => collapsed.value.delete(id)
  const toggle   = (id: string|number) =>
    collapsed.value.has(id) ? expand(id) : collapse(id)
  const isCollapsed = (id: string|number) => collapsed.value.has(id)

  return { flatRows, collapse, expand, toggle, isCollapsed }
}
